#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：ga_fjsp_oop_elite.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/20 21:53 
'''
"""
没问题！下面把精英保留（Elitism）加进 GA，并保证收敛曲线单调不增（即 best_hist 始终记录迄今为止最小的 Cmax）。为方便你直接用，我把完整单文件脚本更新好了（在原有 OOP 版本基础上仅改动 GA 的进化部分并新增 elite 参数）。

运行方式：保存为 ga_fjsp_oop_elite.py，执行 python ga_fjsp_oop_elite.py
依赖：numpy（可选 matplotlib 画图）
"""

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GA for Flexible Job Shop Scheduling (FJSP) - OOP + Elitism
----------------------------------------------------------
- 实体：Operation, Job, Machine, FJSPProblem
- 解码器：FJSPDecoder
- 遗传算法：PPX + 均匀交叉 + 变异 + 精英保留(Elitism)
- 收敛曲线：单调不增（记录历史最优）

Author: you
"""

import random
import numpy as np
from dataclasses import dataclass
from typing import Dict, List, Tuple, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ========== 基础实体类 ==========

@dataclass
class Operation:
    machine_time: Dict[int, int]  # {machine_id: processing_time}

@dataclass
class Job:
    ops: List[Operation]
    @property
    def n_ops(self) -> int: return len(self.ops)

@dataclass
class Task:
    job_id: int
    op_idx: int
    start: float
    end: float
    @property
    def duration(self) -> float: return self.end - self.start

class Machine:
    def __init__(self, machine_id: int):
        self.id = machine_id
        self.timeline: List[Task] = []
        self.ready_time: float = 0.0
    def schedule(self, job_id: int, op_idx: int, start: float, pt: float):
        st = max(self.ready_time, start)
        en = st + pt
        self.timeline.append(Task(job_id, op_idx, st, en))
        self.ready_time = en
        return st, en

class FJSPProblem:
    def __init__(self, jobs: List[Job], m: int):
        self.jobs = jobs
        self.m = m
    @property
    def n_jobs(self) -> int: return len(self.jobs)
    @property
    def n_ops_total(self) -> int: return sum(job.n_ops for job in self.jobs)
    def feasible_machines(self, j: int, k: int) -> List[int]:
        return list(self.jobs[j].ops[k].machine_time.keys())


# ========== 解码器 ==========

class FJSPDecoder:
    def __init__(self, problem: FJSPProblem):
        self.pb = problem
    def decode(self, op_seq: List[int], mac_seq: List[int]) -> Tuple[float, List[Machine]]:
        job_next = np.zeros(self.pb.n_jobs, dtype=int)
        job_ready = np.zeros(self.pb.n_jobs, dtype=float)
        machines = [Machine(i) for i in range(self.pb.m)]
        mac_ptr = 0
        for j in op_seq:
            k = int(job_next[j])
            if k >= self.pb.jobs[j].n_ops: continue
            feas = self.pb.feasible_machines(j, k)
            if not feas: continue
            idx = mac_seq[mac_ptr] % len(feas)
            mm = int(feas[idx])
            pt = float(self.pb.jobs[j].ops[k].machine_time[mm])
            start = max(job_ready[j], machines[mm].ready_time)
            st, en = machines[mm].schedule(j, k, start, pt)
            job_ready[j] = en
            job_next[j] += 1
            mac_ptr += 1
        makespan = float(np.max(job_ready)) if len(job_ready)>0 else 0.0
        return makespan, machines


# ========== GA（含精英保留 & 单调收敛曲线） ==========

class GAFJSP:
    def __init__(self, problem: FJSPProblem,
                 pop: int = 80, gens: int = 400, pc: float = 0.9, pm: float = 0.2,
                 elite: Optional[int] = None,
                 seed: Optional[int] = 0):
        self.pb = problem
        self.pop = pop
        self.gens = gens
        self.pc = pc
        self.pm = pm
        self.elite = elite if elite is not None else max(2, pop // 10)  # 默认10%精英
        if seed is not None:
            random.seed(seed); np.random.seed(seed)
        self.decoder = FJSPDecoder(problem)
        self.base_seq: List[int] = []
        for j, job in enumerate(self.pb.jobs):
            self.base_seq += [j] * job.n_ops

    # 染色体构造
    def random_individual(self) -> Tuple[List[int], List[int]]:
        op_seq = self.base_seq[:]; random.shuffle(op_seq)
        mac_seq=[]; seen=[0]*self.pb.n_jobs
        for j in op_seq:
            k=seen[j]; feas=self.pb.feasible_machines(j,k)
            mac_seq.append(random.randrange(max(1,len(feas))))
            seen[j]+=1
        return op_seq, mac_seq

    # 交叉
    def crossover_ppx(self, a: List[int], b: List[int]) -> List[int]:
        n=len(a); need={}
        for x in a: need[x]=need.get(x,0)+1
        left=dict(need); res=[]; ai=bi=0
        while len(res)<n:
            pickA=(random.random()<0.5)
            if pickA:
                while ai<n and left.get(a[ai],0)==0: ai+=1
                if ai<n: j=a[ai]; res.append(j); left[j]-=1; ai+=1
                else:
                    while bi<n and left.get(b[bi],0)==0: bi+=1
                    j=b[bi]; res.append(j); left[j]-=1; bi+=1
            else:
                while bi<n and left.get(b[bi],0)==0: bi+=1
                if bi<n: j=b[bi]; res.append(j); left[j]-=1; bi+=1
                else:
                    while ai<n and left.get(a[ai],0)==0: ai+=1
                    j=a[ai]; res.append(j); left[j]-=1; ai+=1
        return res
    def crossover_uniform(self, a: List[int], b: List[int]) -> List[int]:
        return [(aa if random.random()<0.5 else bb) for aa,bb in zip(a,b)]

    # 变异
    def mutate(self, op_seq: List[int], mac_seq: List[int]):
        if len(op_seq)>=2 and random.random()<0.5:
            i,j = random.sample(range(len(op_seq)),2)
            op_seq[i],op_seq[j] = op_seq[j],op_seq[i]
        if random.random()<0.7:
            seen=[0]*self.pb.n_jobs
            for i,j in enumerate(op_seq):
                k=seen[j]; seen[j]+=1
                if random.random()<0.2:
                    feas=self.pb.feasible_machines(j,k)
                    mac_seq[i] = random.randrange(max(1,len(feas)))

    # 适应度
    def fitness(self, ind: Tuple[List[int], List[int]]) -> float:
        op_seq, mac_seq = ind
        mk, _ = self.decoder.decode(op_seq, mac_seq)
        return mk

    # 锦标赛选父
    def tournament_pick(self, pop, fits, k=2):
        best_idx = None
        for _ in range(k):
            idx = random.randrange(len(pop))
            if (best_idx is None) or (fits[idx] < fits[best_idx]):
                best_idx = idx
        return pop[best_idx]

    # 进化（含精英保留；best_hist 单调不增）
    def run(self, verbose: bool = True) -> Tuple[Tuple[List[int], List[int]], float, List[float]]:
        # 初始化
        pop = [self.random_individual() for _ in range(self.pop)]
        fits = [self.fitness(ind) for ind in pop]

        best_hist=[]
        best_so_far = float(min(fits))

        for g in range(1, self.gens+1):
            # 1) 精英保留：拷贝前 elite 个个体
            order = np.argsort(fits)
            elites = [pop[i] for i in order[:self.elite]]
            elite_fits = [fits[i] for i in order[:self.elite]]

            # 2) 产生子代直到填满
            offspring=[]
            while len(offspring) < self.pop - self.elite:
                p1 = self.tournament_pick(pop, fits, k=2)
                p2 = self.tournament_pick(pop, fits, k=2)
                c1=(p1[0][:], p1[1][:]); c2=(p2[0][:], p2[1][:])
                if random.random()<self.pc:
                    c1=(self.crossover_ppx(p1[0],p2[0]),
                        self.crossover_uniform(p1[1],p2[1]))
                    c2=(self.crossover_ppx(p2[0],p1[0]),
                        self.crossover_uniform(p2[1],p1[1]))
                if random.random()<self.pm: self.mutate(*c1)
                if random.random()<self.pm: self.mutate(*c2)
                offspring += [c1, c2]
            offspring = offspring[:self.pop - self.elite]

            # 3) 新种群 = 精英 + 子代
            new_pop = elites + offspring
            new_fits = [self.fitness(ind) for ind in new_pop]

            pop, fits = new_pop, new_fits

            # 4) 更新历史最优，保证单调
            cur_best = float(min(fits))
            best_so_far = min(best_so_far, cur_best)
            best_hist.append(best_so_far)

            if verbose and g % 50 == 0:
                print(f"[GA+Elite] Gen {g:4d}  Best Cmax (so far) = {best_so_far:.2f}")

        # 返回最优
        idx = int(np.argmin(fits))
        return pop[idx], fits[idx], best_hist


# ========== 可视化 ==========

def plot_convergence(best_hist: List[float], title: str = "Convergence (Monotone)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip convergence plot)")
        return
    plt.figure(figsize=(6,4))
    plt.plot(best_hist, lw=2)
    plt.xlabel("Generation"); plt.ylabel("Best-so-far makespan")
    plt.title(title); plt.grid(True, alpha=0.3); plt.tight_layout(); plt.show()

def plot_gantt(machines: List[Machine], title: str = "Gantt Chart"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip Gantt plot)")
        return
    fig, ax = plt.subplots(figsize=(12, 0.8*len(machines)+3))
    yticks, ylabels = [], []
    for idx, mach in enumerate(machines):
        for t in sorted(mach.timeline, key=lambda x: x.start):
            ax.barh(idx, t.duration, left=t.start)
            ax.text(t.start + t.duration/2, idx, f"J{t.job_id}-O{t.op_idx}",
                    ha="center", va="center", fontsize=8, color="white")
        yticks.append(idx); ylabels.append(f"M{mach.id}")
    ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
    ax.set_xlabel("Time"); ax.set_title(title)
    ax.grid(axis="x", alpha=0.3); plt.tight_layout(); plt.show()


# ========== 随机实例构造 ==========

def build_random_fjsp(n_jobs: int = 12, m: int = 6,
                      ops_per_job: int = 4, k_machines_per_op: int = 2,
                      pt_low: int = 1, pt_high: int = 9,
                      seed: Optional[int] = 0) -> FJSPProblem:
    rng = np.random.RandomState(seed)
    jobs: List[Job] = []
    for _ in range(n_jobs):
        ops: List[Operation] = []
        for _o in range(ops_per_job):
            feas = sorted(rng.choice(range(m), size=k_machines_per_op, replace=False))
            mt = {mm: int(rng.randint(pt_low, pt_high+1)) for mm in feas}
            ops.append(Operation(machine_time=mt))
        jobs.append(Job(ops=ops))
    return FJSPProblem(jobs, m)


# ========== Demo 主程序 ==========

def main():
    random.seed(0); np.random.seed(0)
    pb = build_random_fjsp(n_jobs=12, m=6, ops_per_job=4, k_machines_per_op=2, seed=42)

    # 关键：elite 参数可调（默认10%）
    ga = GAFJSP(pb, pop=90, gens=400, pc=0.9, pm=0.25, elite=12, seed=1)
    (best_op_seq, best_mac_seq), best_mk, hist = ga.run(verbose=True)
    print(f"\n[Result] Best makespan (final gen) = {best_mk:.2f}")
    print(f"[Result] Best-so-far (hist end)     = {hist[-1]:.2f}")

    decoder = FJSPDecoder(pb)
    mk, machines = decoder.decode(best_op_seq, best_mac_seq)
    assert abs(mk - best_mk) < 1e-6

    plot_convergence(hist, title="FJSP - GA + Elitism (Best-so-far)")
    plot_gantt(machines, title=f"FJSP GA + Elitism Gantt (Cmax={best_mk:.1f})")

if __name__ == "__main__":
    main()
